# -*- coding: utf-8 -*-
# @Time : 2024-11-2024/11/30
# @File : angle_utils.py

import numpy as np


def calculate_angle_3d(p1, p2, p3):
    """
    计算三维空间中三个点形成的角度
    
    参数:
    - p1, p2, p3: 空间中的三个点，p2为角度顶点
    
    返回:
    - angle: 角度(度)
    """
    a = np.array(p1)
    b = np.array(p2)
    c = np.array(p3)
    
    ba = a - b
    bc = c - b
    
    # 计算点积
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc) + 1e-8)
    # 确保cosine_angle在有效范围内，避免数值误差
    cosine_angle = np.clip(cosine_angle, -1.0, 1.0)
    angle = np.arccos(cosine_angle)
    angle = np.degrees(angle)
    return angle


def estimate_3d_from_2d(keypoints_2d, height=180):
    """
    基于2D关键点估计3D信息
    
    参数:
    - keypoints_2d: 2D关键点坐标 [x, y]
    - height: 假设的人物身高（厘米）
    
    返回:
    - keypoints_3d: 估计的3D关键点坐标 [x, y, z]
    """
    # 确保输入格式正确
    if len(keypoints_2d.shape) == 3:  # 如果是(batch, n_points, 2)格式
        keypoints_2d = keypoints_2d[0]  # 取第一个人的关键点
    
    n_points = len(keypoints_2d)
    
    # 创建3D关键点数组
    keypoints_3d = np.zeros((n_points, 3))
    keypoints_3d[:, :2] = keypoints_2d  # 复制x, y坐标
    
    # 如果检测到的关键点不足，则返回原始坐标加上0深度
    if n_points < 17:  # 假设COCO模型有17个关键点
        print(f"警告: 检测到的关键点数量不足: {n_points}, 无法进行准确的3D估计")
        return keypoints_3d
    
    try:
        # 估计身体尺寸
        # 计算颈部到臀部的长度作为身体比例参考
        neck = (keypoints_2d[5] + keypoints_2d[6]) / 2  # 左右肩的中点作为颈部
        hip = (keypoints_2d[11] + keypoints_2d[12]) / 2  # 左右髋的中点
        torso_length_pixels = np.linalg.norm(neck - hip)
        
        # 根据身高估计比例因子 (身高的大约0.3是躯干长度)
        cm_per_pixel = (height * 0.3) / (torso_length_pixels + 1e-6)  # 避免除零
        
        # 定义参考平面 (假设肩部在同一深度平面)
        ref_depth = 0
        keypoints_3d[5, 2] = ref_depth  # 左肩深度
        keypoints_3d[6, 2] = ref_depth  # 右肩深度
        
        # 估计其他关键点的深度
        # 这里使用了一个简化的方法：根据肢体长度的先验知识估计深度
        # 注意：这是一个粗略估计，不是精确的3D重建
        
        # 手臂深度估计
        # 假设手臂最大长度为身高的0.3倍
        arm_max_length = height * 0.3
        
        # 左臂 (如果索引有效)
        if 5 < n_points and 7 < n_points:
            shoulder_elbow_pixels = np.linalg.norm(keypoints_2d[5] - keypoints_2d[7])  # 左肩-左肘
            shoulder_elbow_cm = shoulder_elbow_pixels * cm_per_pixel
            if shoulder_elbow_cm > arm_max_length * 0.5:  # 如果投影长度超过预期，则估计有深度变化
                # 估计肘部深度 (向前或向后)
                depth_diff = (arm_max_length*0.5)**2 - (shoulder_elbow_cm*0.8)**2
                depth_change = np.sqrt(max(0, depth_diff))
                # 根据肘部y坐标判断是向前还是向后
                if keypoints_2d[7][1] < keypoints_2d[5][1]:  # 肘部高于肩部，可能向前
                    keypoints_3d[7, 2] = ref_depth - depth_change
                else:
                    keypoints_3d[7, 2] = ref_depth + depth_change
                
        # 左手腕 (如果索引有效)
        if 7 < n_points and 9 < n_points:
            elbow_wrist_pixels = np.linalg.norm(keypoints_2d[7] - keypoints_2d[9])
            elbow_wrist_cm = elbow_wrist_pixels * cm_per_pixel
            if elbow_wrist_cm > arm_max_length * 0.5:
                depth_diff = (arm_max_length*0.5)**2 - (elbow_wrist_cm*0.8)**2
                depth_change = np.sqrt(max(0, depth_diff))
                keypoints_3d[9, 2] = keypoints_3d[7, 2] + (0 if depth_change < 0 else depth_change)
        
        # 右臂 (类似左臂)
        if 6 < n_points and 8 < n_points:
            shoulder_elbow_pixels = np.linalg.norm(keypoints_2d[6] - keypoints_2d[8])
            shoulder_elbow_cm = shoulder_elbow_pixels * cm_per_pixel
            if shoulder_elbow_cm > arm_max_length * 0.5:
                depth_diff = (arm_max_length*0.5)**2 - (shoulder_elbow_cm*0.8)**2
                depth_change = np.sqrt(max(0, depth_diff))
                if keypoints_2d[8][1] < keypoints_2d[6][1]:
                    keypoints_3d[8, 2] = ref_depth - depth_change
                else:
                    keypoints_3d[8, 2] = ref_depth + depth_change
                
        # 右手腕
        if 8 < n_points and 10 < n_points:
            elbow_wrist_pixels = np.linalg.norm(keypoints_2d[8] - keypoints_2d[10])
            elbow_wrist_cm = elbow_wrist_pixels * cm_per_pixel
            if elbow_wrist_cm > arm_max_length * 0.5:
                depth_diff = (arm_max_length*0.5)**2 - (elbow_wrist_cm*0.8)**2
                depth_change = np.sqrt(max(0, depth_diff))
                keypoints_3d[10, 2] = keypoints_3d[8, 2] + (0 if depth_change < 0 else depth_change)
        
        # 躯干和下肢深度估计
        if 11 < n_points:
            keypoints_3d[11, 2] = ref_depth + 10  # 左髋 (假设比肩部稍后)
        if 12 < n_points:
            keypoints_3d[12, 2] = ref_depth + 10  # 右髋
        
        # 根据膝盖弯曲程度估计深度
        # 左膝
        if 11 < n_points and 13 < n_points:
            hip_knee_pixels = np.linalg.norm(keypoints_2d[11] - keypoints_2d[13])
            hip_knee_cm = hip_knee_pixels * cm_per_pixel
            leg_max_length = height * 0.45 * 0.5  # 腿长约为身高的45%，大腿占一半
            if hip_knee_cm > leg_max_length * 0.9:  # 降低阈值，避免负数
                depth_diff = (leg_max_length)**2 - (hip_knee_cm*0.85)**2
                depth_change = np.sqrt(max(0, depth_diff))
                keypoints_3d[13, 2] = keypoints_3d[11, 2] + depth_change
        
        # 右膝
        if 12 < n_points and 14 < n_points:
            hip_knee_pixels = np.linalg.norm(keypoints_2d[12] - keypoints_2d[14])
            hip_knee_cm = hip_knee_pixels * cm_per_pixel
            if hip_knee_cm > leg_max_length * 0.9:
                depth_diff = (leg_max_length)**2 - (hip_knee_cm*0.85)**2
                depth_change = np.sqrt(max(0, depth_diff))
                keypoints_3d[14, 2] = keypoints_3d[12, 2] + depth_change
        
        # 踝关节
        if 13 < n_points and 15 < n_points:
            keypoints_3d[15, 2] = keypoints_3d[13, 2] + 5  # 左踝
        if 14 < n_points and 16 < n_points:
            keypoints_3d[16, 2] = keypoints_3d[14, 2] + 5  # 右踝
            
        # 头部关键点
        if 0 < n_points:  # 鼻子
            keypoints_3d[0, 2] = ref_depth - 15  # 鼻子前置
        if 1 < n_points and 2 < n_points:  # 眼睛
            keypoints_3d[1, 2] = ref_depth - 12  # 左眼
            keypoints_3d[2, 2] = ref_depth - 12  # 右眼
        if 3 < n_points and 4 < n_points:  # 耳朵
            keypoints_3d[3, 2] = ref_depth - 8   # 左耳
            keypoints_3d[4, 2] = ref_depth - 8   # 右耳
    
    except Exception as e:
        print(f"3D估计过程中出现错误: {e}")
        # 出错时，返回简单的z=0的3D点
    
    return keypoints_3d


def get_joint_angles_config():
    """
    获取关节角度配置
    
    返回:
    - 关节角度配置列表
    """
    # [0:鼻子, 1:左眼, 2:右眼, 3:左耳, 4:右耳, 5:左肩, 6:右肩, 7:左肘, 8:右肘, 9:左腕, 10:右腕, 11:左髋, 12:右髋, 13:左膝, 14:右膝, 15:左踝, 16:右踝]
    return [
        # 上肢角度 - 跑步/运动中的手臂弯曲和摆动
        {"name": "rightElbowAngle", "points": [6, 8, 10], "color": (0, 255, 0)},      # 右肩-右肘-右腕
        {"name": "leftElbowAngle", "points": [5, 7, 9], "color": (0, 255, 0)},       # 左肩-左肘-左腕
        {"name": "rightUpperArmAngle", "points": [12, 6, 8], "color": (0, 200, 0)},    # 右髋-右肩-右肘
        {"name": "leftUpperArmAngle", "points": [11, 5, 7], "color": (0, 200, 0)},    # 左髋-左肩-左肘
        
        # 下肢角度 - 膝盖弯曲和腿部摆动
        {"name": "rightKneeAngle", "points": [12, 14, 16], "color": (255, 0, 0)},    # 右髋-右膝-右踝
        {"name": "leftKneeAngle", "points": [11, 13, 15], "color": (255, 0, 0)},    # 左髋-左膝-左踝
        {"name": "rightThighAngle", "points": [6, 12, 14], "color": (200, 0, 0)},   # 右肩-右髋-右膝
        {"name": "leftThighAngle", "points": [5, 11, 13], "color": (200, 0, 0)},   # 左肩-左髋-左膝
        
        # 躯干角度 - 跑步姿势的前倾和侧倾
        {"name": "trunkForwardLeanAngle", "points": [6, 12, 15], "color": (0, 0, 255)}, # 右肩-右髋-左踝
        {"name": "trunkLateralBendAngle", "points": [5, 0, 6], "color": (0, 0, 200)},   # 左肩-鼻子-右肩
        {"name": "hipAngle", "points": [5, 11, 12], "color": (0, 0, 150)},  # 左肩-左髋-右髋
        
        # 头部姿势
        {"name": "headForwardTiltAngle", "points": [0, 5, 6], "color": (128, 0, 128)},  # 鼻子-左肩-右肩
        
        # 专项运动角度
        {"name": "strideAngle", "points": [13, 11, 12], "color": (255, 165, 0)},  # 左膝-左髋-右髋
        {"name": "armSwingAngle", "points": [7, 5, 6], "color": (255, 140, 0)},   # 左肘-左肩-右肩
        
        # 平衡与协调
        {"name": "rightHipKneeAnkleCoordinationAngle", "points": [12, 14, 16], "color": (50, 180, 180)},  # 右髋-右膝-右踝
        {"name": "leftHipKneeAnkleCoordinationAngle", "points": [11, 13, 15], "color": (50, 180, 180)},  # 左髋-左膝-左踝
        
        # 球类运动特殊角度
        {"name": "rightThrowingAngle", "points": [6, 8, 10], "color": (180, 100, 0)},  # 右肩-右肘-右手腕
        {"name": "headShoulderAngle", "points": [0, 5, 7], "color": (180, 0, 100)},  # 鼻子-左肩-左肘
    ]


def get_skeleton_connections():
    """
    获取骨架连接关系
    
    返回:
    - 骨架连接列表
    """
    # [0:鼻子, 1:左眼, 2:右眼, 3:左耳, 4:右耳, 5:左肩, 6:右肩, 7:左肘, 8:右肘, 9:左腕, 10:右腕, 11:左髋, 12:右髋, 13:左膝, 14:右膝, 15:左踝, 16:右踝]
    return [
        (5, 7), (7, 9),    # 左臂
        (6, 8), (8, 10),   # 右臂
        (11, 13), (13, 15), # 左腿
        (12, 14), (14, 16), # 右腿
        (5, 6),     # 肩膀
        (5, 11), (6, 12),  # 躯干
        (11, 12),   # 髋部
        (0, 5), (0, 6),    # 头部到肩膀
        (0, 1), (0, 2),    # 鼻子到眼睛
        (1, 3), (2, 4)     # 眼睛到耳朵
    ] 